package cn.edu.recommender

import org.apache.spark.sql.SparkSession

import java.lang.Thread.sleep


// 定义样例类
case class Product(productId: Int, name: String, imageUrl: String, categories: String, tags: String)

case class Rating(userId: Int, productId: Int, score: Double, timestamp: Int)

case class MongoConfig(uri: String, db: String)


object DataLoader {

  val PRODUCT_DATA_PATH = "recommender/DataLoader/src/main/resources/products.csv"
  val RATING_DATA_PATH = "recommender/DataLoader/src/main/resources/ratings.csv"
  val MONGODB_PRODUCT_COLLECTION = "Product"
  val MONGODB_RATING_COLLECTION = "Rating"

  def main(args: Array[String]): Unit = {
    // 定义用到的配置参数
    val config = Map(
      "spark.cores" -> "local[*]",
      "mongo.uri" -> "mongodb://localhost:27017/recommender",
      "mongo.db" -> "recommender"
    )

    val mongoConfig = MongoConfig(config("mongo.uri"), config("mongo.db"))

    // 创建一个 SparkSession
    val spark = SparkSession.builder().appName("DataLoader").master(config("spark.cores")).getOrCreate()

    // 在对 DataFrame 和 Dataset 进行操作许多操作都需要这个包进行支持
    import spark.implicits._

    // 将 Product 数据集加载进来
    val productRDD = spark.sparkContext.textFile(PRODUCT_DATA_PATH)

    //将 ProdcutRDD 装换为 DataSet
    val productDS = productRDD.map(item => {
      val attr = item.split("\\^")
      Product(attr(0).toInt, attr(1).trim, attr(4).trim, attr(5).trim, attr(6).trim)
    }).toDS()

    // 将数据保存到 MongoDB 中
    productDS
      .write
      .option("uri", mongoConfig.uri)
      .option("collection", MONGODB_PRODUCT_COLLECTION)
      .mode("overwrite")
      .format("com.mongodb.spark.sql")
      .save()


    // 将 Rating 数据保存到 MongoDB 中
    val ratingRDD = spark.sparkContext.textFile(RATING_DATA_PATH)

    val ratingDS = ratingRDD.map(item => {
      val attr = item.split(",")
      Rating(attr(0).toInt, attr(1).toInt, attr(2).toDouble, attr(3).toInt)
    }).toDS()

    ratingDS
      .write
      .option("uri", mongoConfig.uri)
      .option("collection", MONGODB_RATING_COLLECTION)
      .mode("overwrite")
      .format("com.mongodb.spark.sql")
      .save()

    sleep(300000)
    // 关闭 Spark
    spark.stop()

  }

}
